{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c09c5b6751161422",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "![snad_logo_3.png]()\n",
    "\n",
    "# MNIST dataset\n",
    "\n",
    "This notebook gives an example of Active Anomaly Detection with `coniferest` and [MNIST](https://en.wikipedia.org/wiki/MNIST_database) dataset.\n",
    "\n",
    "Developers of `conferest`:\n",
    "- [Matwey Kornilov (MSU)](https://matwey.name)\n",
    "- [Vladimir Korolev](https://www.linkedin.com/in/vladimir-korolev-a4195b86/)\n",
    "- [Konstantin Malanchev (LINCC Frameworks / CMU)](https://homb.it), notebook author"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d253b06-f0eb-422a-9361-4f612740b648",
   "metadata": {},
   "source": [
    "**[Run this NB in Google Colab](https://colab.research.google.com/github/snad-space/coniferest/blob/master/docs/notebooks/mnist.ipynb)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "667d45495a1937d4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:32:37.978864Z",
     "start_time": "2024-05-13T21:32:37.976186Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "## Install and import the required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d83297862828bda3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:32:24.899586Z",
     "start_time": "2024-05-13T21:32:21.789724Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Install packages\n",
    "%pip install coniferest\n",
    "%pip install datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "424af44173f02c22",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:32:41.825459Z",
     "start_time": "2024-05-13T21:32:40.530086Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import datasets\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from coniferest.isoforest import IsolationForest\n",
    "from coniferest.pineforest import PineForest\n",
    "from coniferest.session import Session\n",
    "from coniferest.session.callback import TerminateAfter, prompt_decision_callback, Label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2dd4922808bfdf2",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Download and load the MNIST dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f63177054a31722",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Download data from Hugging Faces with `datasets` library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b54d353b5b1e296e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:32:47.807541Z",
     "start_time": "2024-05-13T21:32:44.291421Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "mnist = datasets.load_dataset(\"mnist\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d64f3cb6c02e79ec",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Load the data into numpy arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b663431e29a7a250",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:32:50.932628Z",
     "start_time": "2024-05-13T21:32:47.808743Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "images_train = np.asarray(mnist['train']['image'])\n",
    "images_test = np.asarray(mnist['test']['image'])\n",
    "digits_train = np.asarray(mnist['train']['label'])\n",
    "digits_test = np.asarray(mnist['test']['label'])\n",
    "\n",
    "images = np.concatenate([images_train, images_test])\n",
    "digits = np.concatenate([digits_train, digits_test])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4456d1bd107b2666",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Plot some examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9118c00dff72504",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:32:53.961436Z",
     "start_time": "2024-05-13T21:32:53.681108Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 5, figsize=(10, 5))\n",
    "for i in range(10):\n",
    "    ax[i//5, i%5].imshow(images[digits == i][0], cmap='gray')\n",
    "    ax[i//5, i%5].set_title(f'Digit {i}')\n",
    "    ax[i//5, i%5].axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b61fc3ddf5accad",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Preprocess the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae46720ef1c05ef9",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Select the data to use:\n",
    "- `image` : the original images\n",
    "- `fft` : the power spectrum of the images\n",
    "- `both` : the original images and the power spectrum together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa5f1619f4c31db6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:32:59.368335Z",
     "start_time": "2024-05-13T21:32:59.366297Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "DATA = 'both'  # 'image', 'fft', 'both'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4191ed76648104f4",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Make 2-d FFT of the images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f5a8d5f85d1471f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:33:01.656573Z",
     "start_time": "2024-05-13T21:33:00.478893Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Make 2-d FFT of the images\n",
    "data_fft = np.fft.fft2(images)\n",
    "# Get power spectrum\n",
    "power_spectrum = np.square(np.abs(data_fft))\n",
    "# Normalize the power spectrum by zero frequency\n",
    "power_spectrum = power_spectrum / power_spectrum[:, 0, 0][:, None, None]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85c46dc73efbf661",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Plot some examples of power spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4247275995f49b4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:33:03.212083Z",
     "start_time": "2024-05-13T21:33:02.866016Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 5, figsize=(10, 5))\n",
    "for i in range(10):\n",
    "    ax[i//5, i%5].imshow(np.log(power_spectrum[digits == i][0]), cmap='gray')\n",
    "    ax[i//5, i%5].set_title(f'Digit {i}')\n",
    "    ax[i//5, i%5].axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d609b5f5c58c34a2",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Concatenate images and power spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c69336311bdfd4ac",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:33:07.081717Z",
     "start_time": "2024-05-13T21:33:06.659186Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "if DATA == 'image':\n",
    "    final = np.asarray(images, dtype=np.float32)\n",
    "elif DATA == 'fft':\n",
    "    final = np.asarray(power_spectrum.reshape(-1, 28 * 28), dtype=np.float32)\n",
    "elif DATA == 'both':\n",
    "    final = np.concatenate([images.reshape(-1, 28 * 28), power_spectrum.reshape(-1, 28 * 28)], axis=1)\n",
    "else:\n",
    "    raise ValueError(f\"Unknown value for DATA: {DATA}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b3583c32f326df1",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Classic anomaly detection with Isolation forest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95e784345188fd25",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:33:13.934120Z",
     "start_time": "2024-05-13T21:33:08.612649Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "model = IsolationForest(random_seed=10, n_trees=1000)\n",
    "model.fit(np.array(final))\n",
    "scores = model.score_samples(np.array(final))\n",
    "ordered_index = np.argsort(scores)\n",
    "ordered_digits = digits[ordered_index]\n",
    "\n",
    "print(f\"Top 10 weirdest digits : {ordered_digits[:10]}\")\n",
    "print(f\"Top 10 most normal digits : {ordered_digits[-10:]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c174654444043443",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Plot the top 10 weirdest digits and the top 10 most normal digits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cfdf217fc3b5644",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:33:14.340511Z",
     "start_time": "2024-05-13T21:33:13.935803Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 10, figsize=(20, 5))\n",
    "for i in range(10):\n",
    "    ax[0, i].imshow(images[ordered_index[i]], cmap='gray')\n",
    "    ax[0, i].set_title(f'Digit {ordered_digits[i]}')\n",
    "    ax[0, i].axis('off')\n",
    "    ax[1, i].imshow(images[ordered_index[-i - 1]], cmap='gray')\n",
    "    ax[1, i].set_title(f'Digit {ordered_digits[-i - 1]}')\n",
    "    ax[1, i].axis('off')\n",
    "fig.text(0.1, 0.9, 'Top 10 weirdest digits', ha='left', va='center', fontsize=16)\n",
    "fig.text(0.1, 0.5, 'Top 10 most normal digits', ha='left', va='center', fontsize=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac58e1636a9766c",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Anomaly detection with PineForest"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23c4c537793d1bc",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Set expert budget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5dbc9fb2aa6fc4c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:33:15.918132Z",
     "start_time": "2024-05-13T21:33:15.916185Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "EXPERT_BUDGET = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b86dff439131013",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "First, we need a function which would show us an image, its label and ask us if it is an anomaly.\n",
    "\n",
    "Let's say that even numbers are anomalies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9f10fbe6aa4bc1a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:33:16.825580Z",
     "start_time": "2024-05-13T21:33:16.823135Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def decision(index, x, session):\n",
    "    digit, image = digits[index], images[index]\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(2, 2))\n",
    "    ax.imshow(image, cmap='gray')\n",
    "    ax.set_title(f'Digit {digit}')\n",
    "    ax.axis('off')\n",
    "    plt.show()\n",
    "    \n",
    "    ### UNCOMMENT TO MAKE IT INTERACTIVE\n",
    "    # return prompt_decision_callback(index, x, session)\n",
    "\n",
    "    # Non-interactive\n",
    "    return Label.ANOMALY if digit % 2 == 0 else Label.REGULAR"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc07693999b58c31",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Create a model and a session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe34d7b027260913",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:34:02.453945Z",
     "start_time": "2024-05-13T21:33:17.612897Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "model = PineForest(\n",
    "    # Number of trees to use for predictions\n",
    "    n_trees=256,\n",
    "    # Number of new tree to grow for each decision\n",
    "    n_spare_trees=768,\n",
    "    # Fix random seed for reproducibility\n",
    "    random_seed=0,\n",
    ")\n",
    "session = Session(\n",
    "    data=final,\n",
    "    metadata=np.arange(len(final)),\n",
    "    model=model,\n",
    "    decision_callback=decision,\n",
    "    on_decision_callbacks=[\n",
    "        TerminateAfter(EXPERT_BUDGET),\n",
    "    ],\n",
    ")\n",
    "session.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "539178364d28e40c",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Let's see what we have selected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56539dba0a35fd1c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:34:02.457700Z",
     "start_time": "2024-05-13T21:34:02.455223Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "n_anomalies = len(session.known_anomalies)\n",
    "n_total = len(session.known_labels)\n",
    "print(f\"Anomalies: {n_anomalies}/{n_total} ({n_anomalies/n_total:.2%})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "689b590a41bd81c3",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Let's do the opposite: odd numbers are anomalies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4b90bf3afaa6c52",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-13T21:34:45.765519Z",
     "start_time": "2024-05-13T21:34:02.458384Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def decision(index, x, session):\n",
    "    digit, image = digits[index], images[index]\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(2, 2))\n",
    "    ax.imshow(image, cmap='gray')\n",
    "    ax.set_title(f'Digit {digit}')\n",
    "    ax.axis('off')\n",
    "    plt.show()\n",
    "    \n",
    "    ### UNCOMMENT TO MAKE IT INTERACTIVE\n",
    "    # return prompt_decision_callback(index, x, session)\n",
    "\n",
    "    # Non-interactive\n",
    "    return Label.ANOMALY if digit % 2 == 1 else Label.REGULAR\n",
    "\n",
    "model = PineForest(\n",
    "    # Number of trees to use for predictions\n",
    "    n_trees=256,\n",
    "    # Number of new tree to grow for each decision\n",
    "    n_spare_trees=768,\n",
    "    # Fix random seed for reproducibility\n",
    "    random_seed=0,\n",
    ")\n",
    "session = Session(\n",
    "    data=final,\n",
    "    metadata=np.arange(len(final)),\n",
    "    model=model,\n",
    "    decision_callback=decision,\n",
    "    on_decision_callbacks=[\n",
    "        TerminateAfter(EXPERT_BUDGET),\n",
    "    ],\n",
    ")\n",
    "session.run()\n",
    "\n",
    "n_anomalies = len(session.known_anomalies)\n",
    "n_total = len(session.known_labels)\n",
    "print(f\"Anomalies: {n_anomalies}/{n_total} ({n_anomalies/n_total:.2%})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6160aeb6-22c6-428a-a94e-545295fe3e69",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Change `decision` function to make it interactive and try your own experiments. For example, say yes to weird sevens only"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
